cmake_minimum_required(VERSION 3.16 FATAL_ERROR)
project(cuda_kernels LANGUAGES CXX CUDA)

include(CMakePrintHelpers)

# Avoid warning about DOWNLOAD_EXTRACT_TIMESTAMP in CMake >=3.24:
if (CMAKE_VERSION VERSION_GREATER_EQUAL "3.24.0")
    cmake_policy(SET CMP0135 NEW)
endif()

set(CMAKE_EXPORT_COMPILE_COMMANDS on)
set(CMAKE_VERBOSE_MAKEFILE ON)


set(CUDA_KERNELS_DIR   "${PROJECT_SOURCE_DIR}")
set(INCLUDE_DIR        "${PROJECT_SOURCE_DIR}/operator")
set(SRC_DIR            "${PROJECT_SOURCE_DIR}/operator")
set(TEST_DIR           "${PROJECT_SOURCE_DIR}/test")

list(APPEND CMAKE_MODULE_PATH ${CUDA_KERNELS_DIR}/cmake)
include(FetchGTest)
include(FetchSpdlog)
include(FetchTorch)
include(FetchWandb)
set(CMAKE_PREFIX_PATH ${TORCH_SOURCE_DIR}/share/cmake/Torch)
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

set(CMAKE_BUILD_TYPE Release)

function(build_cuda_kernel OP_NAME)
    file(GLOB_RECURSE OP_CPP_FILES  "${SRC_DIR}/${OP_NAME}/*.cpp")
    file(GLOB_RECURSE OP_CU_FILES  "${SRC_DIR}/${OP_NAME}/*.cu")
    file(GLOB_RECURSE OP_TEST_FILES  "${TEST_DIR}/${OP_NAME}/*")
    add_executable(test_${OP_NAME} ${OP_TEST_FILES} ${OP_CPP_FILES} ${OP_CU_FILES})
    target_link_libraries(test_${OP_NAME}
        PRIVATE
        spdlog
        gtest_main
        ${TORCH_LIBRARIES}
        wandbcpp
        )
    target_include_directories(test_${OP_NAME} PRIVATE ${INCLUDE_DIR}/${OP_NAME})
    set_target_properties(test_${OP_NAME} PROPERTIES CUDA_ARCHITECTURES "80;90")
    target_compile_features(test_${OP_NAME} PRIVATE cxx_std_20)
    target_compile_features(test_${OP_NAME} PRIVATE cuda_std_20)
endfunction(build_cuda_kernel)

build_cuda_kernel(pointwise)
# build_cuda_kernel(conv_fwd)